import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from matplotlib.ticker import ScalarFormatter

# ================= 参数配置 =================
# 网络基础参数
BASE_LATENCY = 5  # 基础传播延迟(ms)
STATIC_BW = 40  # 静态路由带宽(Gbps)
SEMI_STATIC_BW = 60  # 半静态路由带宽(Gbps)
HYBRID_BW = 100  # 光电协同带宽(Gbps)
TRAFFIC_RANGE = np.logspace(6, 10, 50)  # 流量范围1MB~10GB(单位：字节)


# ================= 延迟计算模型 =================
def static_routing_latency(size_bytes):
    """静态路由延迟模型：固定带宽 + 线性增长"""
    size_gb = size_bytes / (1024 ** 3)
    transmission_time = (size_gb * 8) / STATIC_BW * 1000  # 转换为ms
    congestion = 0.001 * (size_gb ** 1.5)  # 非线性拥塞因子
    return BASE_LATENCY + transmission_time + congestion


def semi_static_latency(size_bytes):
    """半静态路由延迟模型：动态带宽调整 + 路径切换开销"""
    size_gb = size_bytes / (1024 ** 3)
    effective_bw = SEMI_STATIC_BW * (1 - np.exp(-size_gb / 5))  # 带宽自适应
    transmission_time = (size_gb * 8) / effective_bw * 1000
    reconfig_cost = 2 * np.log1p(size_gb)  # 对数增长的重配开销
    return BASE_LATENCY + transmission_time + reconfig_cost


def hybrid_latency(size_bytes):
    """光电协同延迟模型：光路大管道 + 电层动态优化"""
    size_gb = size_bytes / (1024 ** 3)
    optical_portion = 0.8 * (1 - np.exp(-size_gb / 2))  # 光路承载比例
    elec_portion = 1 - optical_portion

    optical_time = (size_gb * optical_portion * 8) / HYBRID_BW * 1000
    elec_time = (size_gb * elec_portion * 8) / (HYBRID_BW * 0.8) * 1000  # 电层带宽利用率为80%
    coordination_cost = 0.5 * np.sqrt(size_gb)  # 协调开销

    return BASE_LATENCY + optical_time + elec_time + coordination_cost


# ================= 生成数据 =================
traffic_gb = TRAFFIC_RANGE / (1024 ** 3)  # 转换为GB单位

latency_static = [static_routing_latency(x) for x in TRAFFIC_RANGE]
latency_semi = [semi_static_latency(x) for x in TRAFFIC_RANGE]
latency_hybrid = [hybrid_latency(x) for x in TRAFFIC_RANGE]

# ================= 可视化设计 =================
plt.figure(figsize=(12, 7))
ax = plt.gca()

# 绘制曲线
line_static, = ax.plot(traffic_gb, latency_static,
                       color='#d62728', linestyle='--', linewidth=2.5)
line_semi, = ax.plot(traffic_gb, latency_semi,
                     color='#2ca02c', linestyle='-.', linewidth=2.5)
line_hybrid, = ax.plot(traffic_gb, latency_hybrid,
                       color='#1f77b4', linestyle='-', linewidth=2.5)

# 坐标轴设置
ax.set_xscale('log')
ax.set_xlabel("Traffic Volume (GB)", fontsize=12, labelpad=10)
ax.set_ylabel("End-to-End Latency (ms)", fontsize=12, labelpad=10)
ax.set_title("Training Completion Latency Comparison\n(Network Bandwidth: Static={}G, Semi={}G, Hybrid={}G)"
             .format(STATIC_BW, SEMI_STATIC_BW, HYBRID_BW),
             fontsize=14, pad=20)

# 刻度格式化
ax.xaxis.set_major_formatter(ticker.FuncFormatter(
    lambda x, _: f"{x:.1f}GB" if x < 10 else f"{int(x)}GB"))
ax.yaxis.set_major_formatter(ticker.ScalarFormatter(useMathText=True))
ax.ticklabel_format(axis='y', style='sci', scilimits=(0, 0))

# 网格与图例
ax.grid(True, which='both', linestyle='--', alpha=0.6)
ax.legend([line_static, line_semi, line_hybrid],
          ['Static Routing', 'Semi-Static Routing', 'Opto-Electronic Co-design'],
          loc='upper left', fontsize=11)

# 关键区域标注
ax.annotate('Optical Advantage Zone',
            xy=(3, latency_hybrid[15]), xytext=(5, 150),
            arrowprops=dict(arrowstyle='->', color='#1f77b4', lw=1.5),
            fontsize=10, color='#1f77b4')

ax.fill_between(traffic_gb, latency_hybrid, latency_semi,
                where=(traffic_gb > 0.5) & (traffic_gb < 5),
                color='green', alpha=0.1, hatch='//')

# ================= 显示图表 =================
plt.tight_layout()
plt.show()